//
//  EncryptedFileURLProtocol.m
//
//  Created by Robin Summerhill on 15/07/2010.
//  Copyright 2010 Aptogo Ltd. All rights reserved.
//
//  Permission is given to use this source code file without charge in any
//  project, commercial or otherwise, entirely at your risk, with the condition
//  that any redistribution (in part or whole) of source code must retain
//  this copyright and permission notice. Attribution in compiled projects is
//  appreciated but not required.
//

#import "EncryptedFileURLProtocol.h"
#import "NSDataEncrypt.h"

// Default 256bit key - change this!
static NSString *sharedKey = @"abcdefghijklmnopqrstuvwxyz123456";

// Buffer length - must be multiple of 16 for AES256
static const int BUFFER_LENGTH = 512;

static NSString * const ERROR_DOMAIN = @"EncryptedFileURLProtocol";

enum ERROR_CODES {
    DECRYPTION_ERROR_CODE = 1
};

@implementation EncryptedFileURLProtocol
{
    NSInputStream *inStream;
    NSMutableData *dataLoaded;
    uint8_t *inBuffer;
}

+ (BOOL)canInitWithRequest:(NSURLRequest *)request
{
    return YES; // return ([[[request URL] scheme] isEqualToString:ENCRYPTED_FILE_SCHEME_NAME]);
}

+ (NSURLRequest *)canonicalRequestForRequest:(NSURLRequest *)request
{
    return request;
}

+ (void)setKey:(NSString *)key
{
    if (key != sharedKey)
    {
        sharedKey = [key copy];
    }
}

+ (NSString*)key
{
    return sharedKey;
}

// Called when URL loading system initiates a request using this protocol. Initialise input stream, buffers and decryption engine.
- (void)startLoading
{
    dataLoaded = [[NSMutableData alloc]init];
    
    inBuffer = malloc(BUFFER_LENGTH);
    inStream = [[NSInputStream alloc] initWithFileAtPath:[[self.request URL] path]];
    [inStream setDelegate:self];
    [inStream scheduleInRunLoop:[NSRunLoop currentRunLoop] forMode:NSDefaultRunLoopMode];
    [inStream open];
}

// Called by URL loading system in response to normal finish, error or abort. Cleans up in each case.
- (void)stopLoading
{
    [inStream close];
    [inStream removeFromRunLoop:[NSRunLoop currentRunLoop] forMode:NSDefaultRunLoopMode];
    inStream = nil;
    dataLoaded = nil;
    free(inBuffer);
}

- (void)stream:(NSStream *)stream handleEvent:(NSStreamEvent)streamEvent
{
    switch(streamEvent) {
        case NSStreamEventHasBytesAvailable:
        {
            size_t len = 0;
            len = [(NSInputStream *)inStream read:inBuffer maxLength:BUFFER_LENGTH];
            if (len)
            {
                [dataLoaded appendBytes:inBuffer length:len];
            }
            break;
        }
        case NSStreamEventEndEncountered:
        {
            NSData* decrypted = nil;
            
            @try
            {
                decrypted = [dataLoaded aesDecrypt:[sharedKey md5Data]];
            }
            @catch (NSException *exception)
            {
                NSLog(@"%@", exception.description);
            }
            
            if (decrypted != nil)
            {
                [self.client URLProtocol:self didLoadData:decrypted];
                [self.client URLProtocolDidFinishLoading:self];
            }
            else
            {
                [self.client URLProtocol:self didFailWithError:[NSError errorWithDomain:ERROR_DOMAIN code:DECRYPTION_ERROR_CODE userInfo:nil]];
            }
            
            break;
        }
        case NSStreamEventErrorOccurred:
        {
            [self.client URLProtocol:self didFailWithError:[inStream streamError]];
            break;
        }
    }
}

@end
